import au.edu.rmit.tzar.api.Parameters;
import au.edu.rmit.tzar.api.RdvException;
import au.edu.rmit.tzar.api.Runner;
import com.google.common.io.Files;
import edu.monash.nimrod.API.*;
import edu.monash.nimrod.task.TaskCreator;

import java.io.File;
import java.io.IOException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.CountDownLatch;
import java.util.logging.Level;
import java.util.logging.Logger;

/**
 * This runner schedules jobs against Nimrod, using the Nimrod Java API. The parameters for the job are written
 * out to a JSON file, which is copied to the root as part of the nimrod job execution.
 */
public class NimrodRunner implements Runner {
  private static final Logger LOG = Logger.getLogger(NimrodRunner.class.getName());
  private static final String EXPERIMENT_NAME = "test";

  private static final NimrodListener nimrodListener;

  static {
    try {
      nimrodListener = new NimrodListener(NimrodAPI.getNimrodAPI());
    } catch (NimrodException e) {
      throw new RuntimeException(e);
    }
  }

  public boolean runModel(File model, File outputPath, String runId, String flagsString, Parameters parameters,
      Logger logger) throws RdvException {
    File experimentDir = new File(System.getProperty("user.home"), ".nimrod/experiments/" + EXPERIMENT_NAME);
    File parameterFile = new File(experimentDir, "parameters." + runId + ".json");
    File projectFile = new File(experimentDir, "project_data." + runId + ".zip");

    try {
      NimrodAPI nimrod = NimrodAPI.getNimrodAPI();
      // TODO(michaell): pass the name of the experiment (ie source repo url + revision number + project name)
      NimrodExperimentAPI experiment = getOrCreateExperiment(nimrod, EXPERIMENT_NAME);
      copyFilesToNimrodRoot(model, runId, parameters, experimentDir, parameterFile, projectFile);

      NimrodJobAPI job = experiment.addJob(null, "ready", createTask(flagsString, projectFile, parameterFile, runId));

      FutureResult result = nimrodListener.addJobCallback(job.getJobName());

      experiment.start();
      return result.get();
    } catch (NimrodException e) {
      throw new RdvException(e);
    } catch (InterruptedException e) {
      Thread.currentThread().interrupt();
      throw new RdvException(e);
    }
  }

  private static void copyFilesToNimrodRoot(File model, String runId, Parameters parameters, File experimentDir,
      File parameterFile, File projectFile) throws RdvException {
    if (parameterFile.exists()) {
      LOG.warning("Parameters file for run: " + runId + " already exists. Overwriting.");
      parameterFile.delete();
    }

    if (projectFile.exists()) {
      LOG.warning("Project file for run: " + runId + " already exists. Overwriting.");
      projectFile.delete();
    }

    try {
//      parameters.toJson(parameterFile);
      CompressUtil.zipDirectory(projectFile, model, new File("R"), new File("projects"));
      // note that the below line requires that rdv.jar be in the current working directory
      Files.copy(new File("rdv.jar"), new File(experimentDir, "rdv.jar"));
    } catch (IOException e) {
      throw new RdvException(e);
    }
  }

  private TaskCreator createTask(String flags, File projectFile, File parametersFile, String runId) {
    TaskCreator taskCreator = new TaskCreator();
    taskCreator.copyRootToNode(parametersFile.getAbsolutePath(), parametersFile.getName());
    taskCreator.copyRootToNode("rdv.jar", "rdv.jar");
    taskCreator.copyRootToNode(projectFile.getAbsolutePath(), projectFile.getName());
    taskCreator.execute("unzip " + projectFile.getName());
    String outputDir = "tzar." + runId;
    taskCreator.execute("${HOME}/bin/java -Drunnerclass=JythonRunner -jar rdv.jar execlocalruns --revision=-1 " +
        "--runspec=" + parametersFile + " --localcodepath=. --tzarbasedir=" + outputDir + " --commandlineflags=\"" +
        flags + " --rlocation=${HOME}/bin/Rscript\" > stdout");
    taskCreator.copyNodeToRoot("stdout", "stdout." + runId);
    taskCreator.copyNodeToRoot(outputDir, outputDir);
    return taskCreator;
  }

  private NimrodExperimentAPI getOrCreateExperiment(NimrodAPI nimrod, String expname) throws NimrodException {
    try {
      return nimrod.getExperiment(expname);
    } catch (NimrodInvalidExperimentException e) {
      NimrodExperimentAPI experiment = nimrod.addNimrodGExperiment(expname, new String[0], new TaskCreator(), "", "G");
//      experiment.addResource(nimrod.getResource("https://w-entgridglobus-01.infotech.monash.edu.au:8443/SGE@"));
      experiment.addResource(nimrod.getResource("fork"));
      return experiment;
    }
  }

  /**
   * Represents a future result of an operation (success or failure).
   * One thread sets the result, while other threads wait on the result.
   */
  private static class FutureResult {
    private volatile boolean result = false;
    private final CountDownLatch latch = new CountDownLatch(1);

    /**
     * Sets the result, and releases the latch indicating that the result is ready.
     *
     * @param result
     */
    public void set(boolean result) {
      this.result = result;
      latch.countDown();
    }

    /**
     * Gets the result, blocking until it is ready.
     *
     * @return the result of the operation
     * @throws InterruptedException if the thread is interrupted while waiting
     */
    public boolean get() throws InterruptedException {
      latch.await();
      return result;
    }
  }

  private static class NimrodListener {
    private final ConcurrentMap<String, FutureResult> callbacks = new ConcurrentHashMap<String, FutureResult>();

    public NimrodListener(NimrodAPI nimrod) throws NimrodException {
      nimrod.addJobListener(new NimrodJobListener() {
        public void nimrodJobChange(NimrodJobNotification[] njns) {
          for (NimrodJobNotification njn : njns) {
            String jobName = njn.job.getJobName();
            LOG.info("Nimrod job " + jobName + " status " + njn.status);
            FutureResult result;
            if ((result = callbacks.get(jobName)) != null) {
              if (njn.status.equals("done")) {
                callbacks.remove(jobName);
                result.set(true);
              } else if (njn.status.equals("failed")) {
                callbacks.remove(jobName);
                result.set(false);
              }
            }
          }
        }

        public void notificationException(Exception e) {
          LOG.log(Level.WARNING, "Notification exception thrown by Nimrod", e);
        }
      });
    }

    public FutureResult addJobCallback(String jobName) {
      FutureResult result = new FutureResult();
      callbacks.put(jobName, result);
      return result;
    }
  }
}
